#include "mex.h"
#include "ml.h"

void mexFunction( int nlhs, mxArray *plhs[], int nrhs, const mxArray*prhs[] )
{ 	
	/*
		Entry point. This function is called when the DLL is used in Matlab.

		Parameters:		
			+ nlhs: Number of output arguments.
			+ plhs: Array of output matrix pointers. 
			+ nrhs: Number of input parameters.
			+ prhs: Array of input matrix pointers.
	*/

	const char* UsageMsg="[mu,sigma,c]=MEXGmmEM(Data,EMPars,[InitMethod,InitPars],[Mo,So,Co])\nInput Parameters:\n\tData: Matrix NxD with N D-dimension samples.\n\tEMPars: Vector with the EM stop parameters [MaxIters,Min EPS] or an empty list to use the default ones.\n\tInitMethod: Initialization Method (0->Naive Bayes Estimation, 1-> K-Means).\n\tInitPars: Parameters for the initialization Method:\n\t\tNBE: Vector 1xM with the NBE parameters (Optional).\n\t\tK-Means: Number of centres\n\tMo: Initial centers.\n\tSo: Initial covariance matrices\n\tCo: Initial center weights\n";
		
	CvEM em_model;
    CvEMParams params;
	int DimSamples,NumSamples;
	double *pData=NULL;
	int InitMethod,NumClusters;
	double *pNBEParameters=NULL;
	CvMat InitMeans,InitWeights;
	CvMat *InitCovars=NULL;	
	register int c;
	int DimsVec[3];
	unsigned char *pMat=NULL;
	int MaxIters=10;
	double MinEPS=0.1;

	// Check the number of parameters
	if(nrhs!=2 && nrhs!=4 && nrhs!=5)
		mexErrMsgTxt(UsageMsg);

	// Verify the data matrix type
	if(!mxIsNumeric(prhs[0]))
		mexErrMsgTxt("Data Matrix must be double or single");

	// Get the data parameters
	NumSamples=mxGetM(prhs[0]);
	DimSamples=mxGetN(prhs[0]);

	// Get the EM parameters
	if(!mxIsEmpty(prhs[1]))
	{
		double *pEMPars=mxGetPr(prhs[1]);
		MaxIters=pEMPars[0];
		MinEPS=pEMPars[1];
	}
	
	// Point to the data matrix
	pData=mxGetPr(prhs[0]);

	// Assign the initialization method
	switch(nrhs)
	{
	case 2: 	
		InitMethod=0;
		break;
	case 4:	
		InitMethod=*(mxGetPr(prhs[2]));	
		break;
	case 5:
		InitMethod=-1;
		break;
	default:
		mexErrMsgTxt(UsageMsg);
	}
	
	// Get the initialization parameters values	
	switch(InitMethod)
	{
	case 0:// NBE initialization		
		if(nrhs!=2 && nrhs!=4)
			mexErrMsgTxt(UsageMsg);
		else
			pNBEParameters=mxGetPr(prhs[3]);
		break;
	case 1:// K-Means initialization
		if(nrhs!=4)
			mexErrMsgTxt(UsageMsg);
		else
			NumClusters=*(mxGetPr(prhs[3]));
		break;
	}

	// Initialize the parameters
	int NBENumClusters;
	const int *DimsCov=NULL;
	int AreaCov;
	switch(InitMethod)
	{
	case -1:// Initial parameters assigned.
		// Obtain the number of centres (clusters)
		NumClusters=mxGetM(prhs[2]);
		cvInitMatHeader(&InitMeans,mxGetM(prhs[2]),mxGetN(prhs[2]),CV_64FC1,mxGetPr(prhs[2]));
		cvInitMatHeader(&InitWeights,mxGetM(prhs[4]),mxGetN(prhs[4]),CV_64FC1,mxGetPr(prhs[4]));

		// Assing the initial covariance matrices
		params.covs=(const CvMat**)new CvMat*[NumClusters];
		DimsCov=(const int*)mxGetDimensions(prhs[3]);
		AreaCov=DimsCov[0]*DimsCov[1];
		for(c=0;c<NumClusters;c++)
		{
			pMat=(unsigned char*)mxGetData(prhs[3]);
			pMat+=sizeof(double)*AreaCov*c;
			params.covs[c]=cvCreateMatHeader(DimsCov[0],DimsCov[1],CV_64FC1);
			cvSetData((CvMat*)(params.covs[c]),pMat,params.covs[c]->step);
		}
		
		// Assign the initial model to the EM parameters		
		params.means     = &InitMeans;
		params.weights   = &InitWeights;		
		params.nclusters = NumClusters;
		params.start_step= CvEM::START_E_STEP;
		params.cov_mat_type= CvEM::COV_MAT_DIAGONAL;
		break;
	case 0:// NBE initialization

		// Use NBE to initialize the model
		NBENumClusters=0;
		
		// Assign the initial model to the EM parameters
		params.covs      = NULL;
		params.means     = NULL;
		params.weights   = NULL;		
		params.nclusters = NBENumClusters;
		params.start_step= CvEM::START_E_STEP;
		params.cov_mat_type= CvEM::COV_MAT_SPHERICAL;
		break;
	case 1:// K-Means initialization
		params.covs      = NULL;
		params.means     = NULL;
		params.weights   = NULL;		
		params.nclusters = NumClusters;
		params.start_step= CvEM::START_AUTO_STEP;
		params.cov_mat_type= CvEM::COV_MAT_SPHERICAL;
		break;
	}

	// Set the model's parameters
	params.probs              = NULL;
    params.term_crit.max_iter = MaxIters;
    params.term_crit.epsilon  = MinEPS;
    params.term_crit.type     = CV_TERMCRIT_ITER|CV_TERMCRIT_EPS;  

	// Create the Data matrix
	CvMat* Samples = NULL;
	if(mxIsDouble(prhs[0]))
	{		
		CvMat* INSamples = cvCreateMatHeader(NumSamples,DimSamples, CV_64FC1);
		Samples = cvCreateMat(NumSamples,DimSamples, CV_32FC1);
		cvSetData(INSamples,pData,INSamples->step);
		cvConvert(INSamples,Samples);
	}
	else
	{
		Samples = cvCreateMatHeader(NumSamples,DimSamples, CV_32FC1);
		cvSetData(Samples,pData,Samples->step);
	}

	// Learn the model
    em_model.train( Samples, 0, params);
	
	// Point to the return data
	const CvMat *Means=em_model.get_means();
	const CvMat **Covars=em_model.get_covs();
	const CvMat *Weights=em_model.get_weights();
	
	switch(cvGetElemType(Means))
	{
	case CV_64FC1:
		plhs[0]=mxCreateNumericMatrix(Means->rows,Means->cols,mxDOUBLE_CLASS,mxREAL);
		memcpy(mxGetData(plhs[0]),Means->data.ptr,sizeof(double)*Means->cols*Means->rows);
		break;
	case CV_32FC1:
		plhs[0]=mxCreateNumericMatrix(Means->rows,Means->cols,mxSINGLE_CLASS,mxREAL);
		memcpy(mxGetData(plhs[0]),Means->data.ptr,sizeof(float)*Means->cols*Means->rows);
		break;
	}

	if(nlhs>1)
	{		
		DimsVec[0]=Covars[0]->rows;
		DimsVec[1]=Covars[0]->cols;
		DimsVec[2]=em_model.get_nclusters();

		switch(cvGetElemType(Covars[0]))
		{
		case CV_64FC1:			
			plhs[1]=mxCreateNumericArray(3,(const mwSize*)DimsVec,mxDOUBLE_CLASS,mxREAL);
			pMat=(unsigned char*)mxGetData(plhs[1]);
			for(c=0;c<em_model.get_nclusters();c++)
			{	
				memcpy(&(pMat[sizeof(double)*Covars[c]->cols*Covars[c]->rows*c]),Covars[c]->data.ptr,sizeof(double)*Covars[c]->cols*Covars[c]->rows);
			}
			break;
		case CV_32FC1:
			plhs[1]=mxCreateNumericArray(3,(const mwSize*)DimsVec,mxSINGLE_CLASS,mxREAL);
			pMat=(unsigned char*)mxGetData(plhs[1]);
			for(c=0;c<em_model.get_nclusters();c++)
			{
				memcpy(&(pMat[sizeof(float)*Covars[c]->cols*Covars[c]->rows*c]),Covars[c]->data.ptr,sizeof(float)*Covars[c]->cols*Covars[c]->rows);
			}
			break;
		}
	}

	if(nlhs>2)
	{
		switch(cvGetElemType(Weights))
		{
		case CV_64FC1:
			plhs[2]=mxCreateNumericMatrix(Weights->rows,Weights->cols,mxDOUBLE_CLASS,mxREAL);
			memcpy(mxGetData(plhs[2]),Weights->data.ptr,sizeof(double)*Weights->cols*Weights->rows);
			break;
		case CV_32FC1:
			plhs[2]=mxCreateNumericMatrix(Weights->rows,Weights->cols,mxSINGLE_CLASS,mxREAL);
			memcpy(mxGetData(plhs[2]),Weights->data.ptr,sizeof(float)*Weights->cols*Weights->rows);
			break;
		}
	}	

	// Release temporal data
	if(mxIsDouble(prhs[0]))
		cvReleaseMat(&Samples);
}
